# cli/main.py
import json
from pathlib import Path
import typer

from core.extractor import LineageExtractor
from viz.sankey import render_column_sankey, render_column_to_column_sankey
from viz.tree import render_table_tree

app = typer.Typer(add_completion=False)

@app.command()
def parse(
    file: Path = typer.Argument(..., exists=True, readable=True, help="SQL 文件"),
    viz: str = typer.Option("all", help="可视化类型：table|column|col2col|all"),
    outdir: Path = typer.Option(Path("."), help="输出目录（保存 HTML）"),
    topn: int = typer.Option(20, help="表→列图每表最多展示的字段数"),
):
    sql = file.read_text(encoding="utf-8")
    res = LineageExtractor().extract(sql)

    # 输出 JSON 到控制台（与之前行为一致）
    print(json.dumps({
        "statements": [s.__dict__ | {
            # 将 set 转 list 以便序列化
            "tables": sorted(list(s.tables)),
            "target_tables": sorted(list(s.target_tables)),
            "source_tables": sorted(list(s.source_tables)),
            "columns_by_table": {k: sorted(list(v)) for k, v in s.columns_by_table.items()},
            "ctes": {
                k: (list(v) if isinstance(v, set) else v)
                for k, v in s.ctes.items()
            },
            "column_lineage": {k: sorted(list(v)) for k, v in s.column_lineage.items()},
        } for s in res.statements]
    }, ensure_ascii=False, indent=2))

    # 可视化导出
    outdir.mkdir(parents=True, exist_ok=True)
    for idx, st in enumerate(res.statements):
        base = file.stem if len(res.statements) == 1 else f"{file.stem}_{idx+1}"
        if viz in ("column", "all"):
            html = render_column_sankey(st, top_n_columns=topn)
            (outdir / f"{base}_column_sankey.html").write_text(html, encoding="utf-8")
        if viz in ("table", "all"):
            html = render_table_tree(st)
            (outdir / f"{base}_table_tree.html").write_text(html, encoding="utf-8")
        if viz in ("col2col", "all"):
            html = render_column_to_column_sankey(st)
            (outdir / f"{base}_col2col_sankey.html").write_text(html, encoding="utf-8")

if __name__ == "__main__":
    app()
